package cn.lzd.mr.mroverride.inputformat;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.input.FileSplit;

import java.io.IOException;

/**
 * Created by lzd on 2018/6/28.
 */
public class CustomRecordReader extends RecordReader<IntWritable, Text> {

    private IntWritable intWritable = new IntWritable();
    private Text text = new Text();
    private FileSplit fileSplit;
    private Configuration configuration;
    private boolean flag = false;

    @Override
    public void initialize(InputSplit inputSplit, TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
        this.fileSplit = (FileSplit) inputSplit;
        this.configuration = taskAttemptContext.getConfiguration();
    }

    /**
     * 读取整个文件的内容封装到value中
     * @return
     * @throws IOException
     * @throws InterruptedException
     */
    @Override
    public boolean nextKeyValue() throws IOException, InterruptedException {
        if (!flag) {
            FSDataInputStream in = null;
            try {
                Path path = fileSplit.getPath();
                FileSystem fileSystem = path.getFileSystem(this.configuration);
                in = fileSystem.open(path);
                byte[] bytes = new byte[(int) fileSplit.getLength()];
                IOUtils.readFully(in, bytes, 0, bytes.length);
                text.set(bytes, 0, bytes.length);
                flag = true;
            }finally {
                IOUtils.closeStream(in);
            }
            return flag;
        }
        return false;
    }

    @Override
    public IntWritable getCurrentKey() throws IOException, InterruptedException {
        return this.intWritable;
    }

    @Override
    public Text getCurrentValue() throws IOException, InterruptedException {
        return this.text;
    }

    @Override
    public float getProgress() throws IOException, InterruptedException {
        return this.flag ? 1: 0;
    }

    @Override
    public void close() throws IOException {

    }
}
